import time
import os
import tkinter as tk
from PIL import ImageGrab
from pynput import keyboard

# --- 配置 ---
SAVE_DIR = "D:\\fk2025"
ACTIVATION_KEY = keyboard.KeyCode.from_char('q') 
# 确保文件夹存在
os.makedirs(SAVE_DIR, exist_ok=True)

# --- 全局状态 ---
app = None # Tkinter App 实例
start_x, start_y = 0, 0
current_x, current_y = 0, 0
is_capturing = False # 是否处于鼠标按下拖动状态

# --- 截图类：负责图形绘制和截图逻辑 ---
class ScreenshotApp(tk.Toplevel):
    def __init__(self, master=None):
        super().__init__(master)
        
        # 1. 窗口配置：实现绿色透明和置顶
        self.attributes('-alpha', 0.2)      # <--- 提高透明度到 0.2，让绿色可见
        self.attributes('-fullscreen', True) # 全屏覆盖
        self.attributes('-topmost', True)   # 保持在最上层
        self.overrideredirect(True)         # 移除窗口边框和标题栏
        self.withdraw()                     # 初始时隐藏窗口

        # 2. Canvas 配置：用于绘制边框
        # 设置 Canvas 背景色为纯绿色，结合窗口透明度，形成透明绿色覆盖
        self.canvas = tk.Canvas(self, cursor="cross", highlightthickness=0, bg='#00FF00') 
        self.canvas.pack(fill=tk.BOTH, expand=tk.YES)
        
        # 3. 鼠标事件绑定
        self.canvas.bind("<ButtonPress-1>", self.on_mouse_down)
        self.canvas.bind("<B1-Motion>", self.on_mouse_move) # 鼠标左键按下并移动
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_up)
        
        # 4. 实时坐标标签 (用于显示鼠标当前位置)
        # 标签颜色使用亮黄色，保持显眼
        self.coord_label = tk.Label(self.canvas, text="", bg="black", fg="yellow", font=("Arial", 10))
        self.coord_label.place(x=10, y=10)
        
        # 5. 键盘事件绑定 (按 Esc 退出截图模式)
        self.bind('<Escape>', self.cancel_capture)

    def activate_capture_mode(self):
        """激活截图模式：显示全屏透明窗口"""
        self.deiconify() # 显示窗口
        self.focus_force() # 强制获取焦点
        print("\n--- 截图模式已激活，请拖动鼠标选择区域或按 Esc 取消 ---")

    def cancel_capture(self, event=None):
        """取消截图操作：隐藏窗口"""
        global is_capturing
        self.withdraw() # 隐藏窗口
        is_capturing = False
        print("--- 截图模式已取消 ---")

    def on_mouse_down(self, event):
        """鼠标左键按下：记录起始点"""
        global start_x, start_y, is_capturing
        start_x, start_y = event.x_root, event.y_root
        is_capturing = True
        self.canvas.delete("rect") # 清除上次绘制的边框
        
    def on_mouse_move(self, event):
        """鼠标左键拖动：实时绘制边框和更新坐标"""
        if is_capturing:
            global current_x, current_y
            current_x, current_y = event.x_root, event.y_root
            
            # 绘制矩形边框（亮黄色，4px）
            self.canvas.delete("rect")
            self.canvas.create_rectangle(
                start_x, start_y, current_x, current_y,
                outline="yellow", # 边框颜色改为亮黄色
                width=4,          # 边框宽度增加到 4 像素
                tags="rect"
            )
            
            # 更新鼠标位置显示
            width = abs(current_x - start_x)
            height = abs(current_y - start_y)
            self.coord_label.config(text=f"位置: ({current_x}, {current_y})\n大小: {width}x{height}")
            # 移动标签到鼠标附近
            self.coord_label.place(x=current_x + 15, y=current_y + 15)


    def on_mouse_up(self, event):
        """鼠标左键释放：执行截图并隐藏窗口"""
        global is_capturing
        
        if is_capturing:
            end_x, end_y = event.x_root, event.y_root
            self.perform_screenshot(start_x, start_y, end_x, end_y)
        
        is_capturing = False

    def perform_screenshot(self, x1, y1, x2, y2):
        """根据坐标截图并保存文件，确保在截图前隐藏边框"""
        
        # 1. 关键步骤：在抓取屏幕内容前，瞬间隐藏包含边框的窗口
        self.withdraw() 
        
        # 2. 规范坐标
        left = min(x1, x2)
        top = min(y1, y2)
        right = max(x1, x2)
        bottom = max(y1, y2)
        
        save_path = None
        try:
            # 3. 抓取屏幕截图 (抓取的是截图框下面的实际内容)
            img = ImageGrab.grab(bbox=(left, top, right, bottom))
            
            # 构造文件名: test时间戳.jpg
            timestamp = time.strftime("%Y%m%d%H%M%S")
            file_name = f"test{timestamp}.jpg"
            save_path = os.path.join(SAVE_DIR, file_name)
            
            img.save(save_path, "JPEG")
            print(f"✅ 截图成功并保存到: {save_path}")
            
            # 返回保存的文件路径，供下一步 OCR 使用
            return save_path
            
        except Exception as e:
            print(f"❌ 截图失败: {e}")
            return None


# --- 全局键盘监听器 (用于激活截图模式) ---
def on_press(key):
    """监听键盘按下事件"""
    global app, is_capturing

    # 激活截图模式的按键
    if key == ACTIVATION_KEY and not is_capturing:
        if app:
            app.activate_capture_mode()
        return

    # 按 Esc 键退出程序
    if key == keyboard.Key.esc:
        if app and app.winfo_ismapped(): # 如果截图窗口是显示的，先取消截图
            app.cancel_capture()
        else: # 否则退出整个程序
            print("\n--- 检测到 Esc 键，程序退出 ---")
            # 停止 Tkinter 主循环
            if app:
                app.master.quit()
            return False # 停止 pynput 监听器

# --- 主程序启动 ---
if __name__ == "__main__":
    # 1. 启动 Tkinter 主窗口
    root = tk.Tk()
    root.withdraw() # 隐藏根窗口
    
    # 2. 实例化截图应用
    app = ScreenshotApp(root)

    # 3. 启动 pynput 键盘监听 (在单独线程中)
    k_listener = keyboard.Listener(on_press=on_press)
    k_listener.daemon = True # 设置为守护线程
    k_listener.start()

    print("**********************************************")
    print(f"🚀 截图程序启动，保存目录: {SAVE_DIR}")
    print(f"📌 按下键 {ACTIVATION_KEY.char.upper()} 激活截图模式")
    print("📌 激活后，按住鼠标左键拖动以绘制边框并截图")
    print("📌 随时按 Esc 键退出程序")
    print("**********************************************")

    # 4. 启动 Tkinter 主循环 (GUI 必须在主线程运行)
    root.mainloop()

    # 5. 清理 pynput 监听器
    k_listener.stop()
    print("程序已安全终止。")